/**
 * 
 */
package cn.edu.bjtu.textprocess.core;

import java.util.ArrayList;
import java.util.List;

import org.deeplearning4j.iterator.LabeledSentenceProvider;
import org.deeplearning4j.iterator.provider.CollectionLabeledSentenceProvider;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.textplantform.common.core.ClassificationPair;
import org.textplantform.common.provider.ClassificationService;

import cn.edu.bjtu.model.core.Deep4jNetworkWrapper;
import cn.edu.bjtu.model.core.TextCategorizationManager;
import cn.edu.bjtu.util.DatasetUtils;


/**
 * @author alex
 *
 */
public class TextClassificationModel implements ClassificationService{
	int maxSentenceLength = 256;
	
	Deep4jNetworkWrapper net;
	Word2VecModel w2v;

	public TextClassificationModel(Word2VecModel wv) throws Exception {
		w2v = wv;
		net =new Deep4jNetworkWrapper(new LocalFileModelReader(),-1);
		net.load();
	}
	
	@Override
	public ClassificationPair[] classifyDocument(String doc) throws Exception {
		return null;
	}

	@Override
	public ClassificationPair[][] classifyDocument(String[] docs) throws Exception {
		return ff(docs);
	}

	private ClassificationPair[][] ff(String[] docs) throws Exception{
		List<String> docList = new ArrayList<String>(docs.length);
		List<String> labelList = new ArrayList<>(docs.length);
		for(int i=0;i<docs.length;i++){
			docList.add(docs[i]);
			labelList.add("unknow");
		}
		//第三个参数一定要有,防止洗牌
		LabeledSentenceProvider lsp = new CollectionLabeledSentenceProvider(docList,labelList,null);
		List<INDArray[]> resList = new ArrayList<INDArray[]>(5);
		try{
			DataSetIterator dsi = DatasetUtils.getCNNDataSet(lsp, 32, maxSentenceLength, w2v.getW2v());
			while(dsi.hasNext()){
				DataSet ds = dsi.next();
				resList.add(net.output(ds.getFeatures()));
			}
		}finally{
		}
		//减少同步块.
		ClassificationPair[][] result = new ClassificationPair[docs.length][];
		int k = 0;
		for(int i=0;i<resList.size();i++){
			INDArray ele = resList.get(i)[0];
			for(int j = 0;j<ele.rows();j++){
				result[k++] = TextCategorizationManager.get().getDesc(ele.getRow(j));
			}
		}
		return result;
	}
	

}
